#!/usr/bin/python
# -*- coding: utf-8 -*-

import numpy as np
import pandas as pd
import time

SPEC_NUM = 5
ROW_NUM = SPEC_NUM
COL_NUM = SPEC_NUM
N_STATES = 847288609443   # 1维世界的宽度
ACTIONS = [_ for _ in range(ROW_NUM * COL_NUM)]     # 探索者的可用动作
EPSILON = 0.9   # 贪婪度 greedy
ALPHA = 0.1     # 学习率
GAMMA = 0.9    # 奖励递减值
MAX_EPISODES = 13  # 最大回合数
FRESH_TIME = 0.3    # 移动间隔时间

def build_q_table(n_states, actions):
    table = pd.DataFrame(
        np.zeros((n_states, len(actions))),     # q_table 全 0 初始
        columns=actions,    # columns 对应的是行为名称
    )
    return table

def choose_action(state, q_table):
    state_actions = q_table.iloc[state, :]  # 选出这个 state 的所有 action 值
    if (np.random.uniform() > EPSILON) or (state_actions.all() == 0):  # 非贪婪 or 或者这个 state 还没有探索过
        actions = [_ for _ in ACTIONS if get_env_block_state(state, _) == 0] # 筛选出可以下的点
        action_name = np.random.choice(actions)
    else:
        action_name = q_table.columns[state_actions.argmax()]    # 贪婪模式
    return action_name

def get_env_feedback(S, A):
    # This is how agent will interact with the environment
    t = get_env_block_state(S, ROW_NUM * COL_NUM)
    S_ = S + t * math.pow(3, A)
    if is_terminated(S, A):
        R = 1
    else:
        R = 0
    return S_, R

def get_env_block_state(S, pos):
    if (pos >= 0 and pos <= ROW_NUM * COL_NUM):
        return (S // math.pow(3, pos)) % 3
    return -1

def print_state(S):
    print("\033[%d;%df" % (18, 0), end="")
    for i in range(ROW_NUM * COL_NUM):
        s = (S // math.pow(3, i)) % 3
        s = " " if (s == 0) else "X" if (s == 1) else "O"
        e = '' if ((i + 1) % ROW_NUM != 0) else '\n'
        print(f" {s}", end=e)
    return

def is_terminated(S, A):
    bstate = get_env_block_state(S, A)
    for i in (ROW_NUM + 1, ROW_NUM, ROW_NUM - 1, 1):
        steps = 1
        for off in (-i, +i):
            cur = A + off
            while (get_env_block_state(S, cur) == bstate):
                steps += 1
                cur += off
        if steps == 5:
            return True
    return False

def update_env(S, A, episode, step_counter):
    # This is how environment be updated
    # env_list = ['-']*(N_STATES-1) + ['T']   # '---------T' our environment
    if is_terminated(S, A):
        interaction = 'Episode %s: total_steps = %s' % (episode+1, step_counter)
        print('\r{}'.format(interaction), end='')
        time.sleep(2)
        print('\r                                ', end='')
    else:
        print_state(S)
        # env_list[S] = 'o'
        # interaction = ''.join(env_list)
        # print('\r{}'.format(interaction), end='')
        time.sleep(FRESH_TIME)

def rl():
    q_table = build_q_table(N_STATES, ACTIONS)  # 初始 q table
    for episode in range(MAX_EPISODES):     # 回合
        step_counter = 0
        S = 0   # 回合初始位置
        A = -1
        is_terminated = False   # 是否回合结束
        update_env(S, A, episode, step_counter)    # 环境更新
        while not is_terminated:

            A = choose_action(S, q_table)   # 选行为
            S_, R = get_env_feedback(S, A)  # 实施行为并得到环境的反馈
            q_predict = q_table.loc[S, A]    # 估算的(状态-行为)值
            if R != 1:
                q_target = R + GAMMA * q_table.iloc[S_, :].max()   #  实际的(状态-行为)值 (回合没结束)
            else:
                q_target = R     #  实际的(状态-行为)值 (回合结束)
                is_terminated = True    # terminate this episode

            q_table.loc[S, A] += ALPHA * (q_target - q_predict)  #  q_table 更新
            S = S_  # 探索者移动到下一个 state

            update_env(S, A, episode, step_counter+1)  # 环境更新

            step_counter += 1
    return q_table



if __name__ == "__main__":
    q_table = rl()
    print("Q-learning")
    print(q_table)
